#!/bin/bash
set -euo pipefail

# 脚本名称：restore_sdcard.sh
# 功能说明：
#   本脚本用于将镜像文件恢复到 SD 卡。
#   主要流程：
#     1. 获取所有可移动设备列表
#     2. 选择要恢复的 SD 卡设备
#     3. 卸载所有已挂载的分区
#     4. 使用 dd 命令将镜像文件恢复到 SD 卡
#     5. 刷新磁盘让系统识别到新的分区表
#
# 注意事项：
#   - 部分命令（如 dd）需要 root 权限，请确保使用 sudo 运行或提前切换为 root。
#   - 执行前请确认数据已备份，避免误操作造成数据丢失。

# 打印信息
echo_info() { echo -e "\e[1;32m[INFO]\e[0m $1"; }
echo_error() { echo -e "\e[1;31m[ERROR]\e[0m $1" >&2; }
error_exit() { echo_error "$1"; exit 1; }

# 1. 获取可移动设备列表（修复 `sed` 问题）
get_sdcard_devices() {
    local devices=()
    for dev in /sys/block/*/removable; do
        [[ -f "$dev" && "$(cat "$dev")" == "1" ]] || continue
        dev_name=$(basename "$(dirname "$dev")")
        devices+=("/dev/$dev_name")
    done
    echo "${devices[@]}"
}

# 2. 选择 SD 卡设备
select_sdcard() {
    local disks=($(get_sdcard_devices))
    [[ ${#disks[@]} -eq 0 ]] && error_exit "未检测到 SD 卡。"

    while true; do
        echo_info "可用的 SD 卡:"
        lsblk -d "${disks[@]}" -o NAME,RM,SIZE,TYPE,MODEL,SERIAL

        read -rp "输入要恢复的设备文件名（或按 Enter 键退出）：/dev/" disk
        [[ -z "$disk" ]] && exit 0

        if [[ -b "/dev/$disk" && " ${disks[*]} " == *" /dev/$disk "* ]]; then
            echo_info "选择的设备：/dev/$disk"
            DISK="/dev/$disk"
            break
        else
            echo_error "无效的设备。请从列表中选择。"
        fi
    done
}

# 3. 卸载所有已挂载的分区
umount_partitions() {
    echo_info "卸载设备 $DISK 上的所有已挂载分区..."
    for PART in $(lsblk "$DISK" -o MOUNTPOINT --noheadings | awk '{print $1}'); do
        sudo umount "$PART" || echo_error "卸载 $PART 失败"
    done
}

# 4. 恢复 SD 卡
restore_sdcard() {
    [[ -f "$1" ]] || error_exit "镜像文件未找到：$1"
    local image_file="$1"

    echo_info "恢复镜像：$image_file -> $DISK"
    read -rp "确认是否继续恢复？（yes/no）： " confirm
    [[ "$confirm" != "yes" ]] && error_exit "恢复已取消。"

    echo_info "开始恢复..."
    sudo dd if="$image_file" of="$DISK" status=progress conv=sync && echo_info "恢复成功完成！"

    echo_info "刷新磁盘以识别新的分区表..."
    sudo partprobe "$DISK" && echo_info "磁盘刷新完成！"
}

# 主流程
main() {
    [[ -z "$1" ]] && error_exit "用法: $0 <镜像文件>"
    select_sdcard
    umount_partitions
    restore_sdcard "$1"
}

main "$@"

